import time


class Trie(object):
    def __init__(self):
        self.search_map = dict()
        self.already_search_dic = dict()

    def insert(self, word: str) -> None:
        dic = self.search_map
        for s in word:
            if s not in dic:
                dic[s] = dict()
            dic = dic[s]
        # self.already_search_dic = dict()

    def startsWith(self, prefix: str) -> bool:
        if prefix not in self.already_search_dic:
            dic = self.search_map
            for s in prefix:
                if s not in dic:
                    self.already_search_dic[prefix] = False
                    break
                else:
                    dic = dic[s]
            else:
                self.already_search_dic[prefix] = True
        return self.already_search_dic[prefix]


class Solution:
    def findWords(self, board: list, words: list) -> list:
        if not board or not words:
            return []
        self.tree, self.res = Trie(), set()
        self.words_set = set(words)
        self.board = board
        self.max_i, self.max_j = len(board) - 1, len(board[0]) - 1
        self.max_len = len(words)
        for word in words:
            self.tree.insert(word)
        try:
            for i, lists in enumerate(board):
                for j, alpha in enumerate(lists):
                    if self.tree.startsWith(alpha):
                        self.search(i, j, {(i, j), }, alpha)
        finally:
            return list(self.res)

    def search(self, i, j, path: set, mid_word: str):
        if mid_word in self.words_set:
            self.res.add(mid_word)
            if len(self.res) == self.max_len:
                raise StopAsyncIteration
        elif not self.tree.startsWith(mid_word):
            return
        if i > 0 and (i - 1, j) not in path and self.tree.startsWith(mid_word + self.board[i-1][j]):
            self.search(i - 1, j, path | {(i - 1, j)}, mid_word + self.board[i - 1][j])
        if i < self.max_i and (i + 1, j) not in path and self.tree.startsWith(mid_word + self.board[i+1][j]):
            self.search(i + 1, j, path | {(i + 1, j)}, mid_word + self.board[i + 1][j])
        if j > 0 and (i, j - 1) not in path and self.tree.startsWith(mid_word + self.board[i][j-1]):
            self.search(i, j - 1, path | {(i, j - 1)}, mid_word + self.board[i][j - 1])
        if j < self.max_j and (i, j + 1) not in path and self.tree.startsWith(mid_word + self.board[i][j+1]):
            self.search(i, j + 1, path | {(i, j + 1)}, mid_word + self.board[i][j + 1])


a = Solution()
print(a.findWords(words=["oath", "pea", "eat", "rain"], board=[['o', 'a', 'a', 'n'],
                                                               ['e', 't', 'a', 'e'],
                                                               ['i', 'h', 'k', 'r'],
                                                               ['i', 'f', 'l', 'v']]))
now = time.time()
print(a.findWords([["b", "b", "a", "a", "b", "a"],
                   ["b", "b", "a", "b", "a", "a"],
                   ["b", "b", "b", "b", "b", "b"],
                   ["a", "a", "a", "b", "a", "a"],
                   ["a", "b", "a", "a", "b", "b"]], ["abbbababaa"]))
print(time.time() - now)
